#! /usr/bin/python3

import argparse, time

def argument():
    arg = argparse.ArgumentParser()
    arg.add_argument("--file", '-f', required = True, help = 'Fasta file name')
    arg.add_argument("--warning", '-w', default = False, type = bool, help = 'Print Warning information')
    return arg.parse_args()

__Warning__ = False


# compute the SP of two seqs
# by cjt
def Compute_two(s1, s2, d=3, e=1, m=1, mis=-2):
    if len(s1) != len(s2):
        # print(s1, s2)
        if __Warning__: print("Warning: the length of s1 and s2 is wrong!")
    maxlen = max(len(s1), len(s2))
    if len(s1) != maxlen: s1 = s1 + '-' * (maxlen - len(s1))
    if len(s2) != maxlen: s2 = s2 + '-' * (maxlen - len(s2))
    score_two = 0
    gap1 = 0
    for i in range(maxlen):
        if s1[i] != "-" and s2[i] != "-":
            if gap1:
               gap1 = 0
            if s1[i] == s2[i]:
                score_two += m
            else:
                score_two += mis
        elif s1[i] != "-" or s2[i] != "-":
            if gap1 == 0:
                score_two -= d
                gap1 = 1
            else:
                score_two -= e
    return score_two


if __name__ == '__main__':
    start_real_time = time.time()
    start_CPU_time = time.perf_counter()
    arg = argument()
    f = open(arg.file)
    __Warning__ = arg.warning
    sequences = []
    thisline = []
    for line in f:
        if line.startswith(">"):
            if thisline != []: sequences.append(''.join(thisline))
            thisline = []
        else:
            thisline.append(line.strip())
    if thisline != []:
        sequences.append(''.join(thisline))
    # print(sequences)
    if len(sequences) < 1: 
        print("Error: this file has no sequences!!")
        exit(1)
    # Check the sequence is the same length
    seqlen1 = len(sequences[0])
    seqlen = [seqlen1]
    maxlen = seqlen1
    sequence_cnt = 1
    for i in range(1, len(sequences)):
        thislen = len(sequences[i])
        if thislen != seqlen1:
            if __Warning__: print("Warning: The length of sequence " + str(i) + "(" + str(thislen) + ") is not same as sequence 0(" + str(seqlen1) + ")!!")
            maxlen = max(maxlen, thislen)
        seqlen.append(thislen)
        sequence_cnt += 1

    Value_SP = 0
    for i in range(len(sequences)):
        for j in range(i + 1, len(sequences)):
            Value_SP += Compute_two(sequences[i], sequences[j])
    
    end_real_time = time.time()
    end_CPU_time = time.perf_counter()
    print("SP Score is: " + str(Value_SP) + "\nUsed time: " + str(end_real_time - start_real_time) + ", CPU time: " + str(end_CPU_time - start_CPU_time))
